package top.chukongxiang.monitorservice.tools.rtp;

import cn.hutool.core.util.ByteUtil;
import lombok.extern.slf4j.Slf4j;
import org.springframework.util.LinkedMultiValueMap;

import java.nio.ByteOrder;
import java.util.List;
import java.util.Optional;

/**
 * PS 流包装
 * @author 楚孔响
 * @version 1.0
 * @date 2023/9/13 19:43
 */
@Slf4j
public class PsWrapper extends StreamWrapper {

    private final LinkedMultiValueMap<RtpPayload.PsPack, byte[]> PS_MAP = new LinkedMultiValueMap<>();

    public PsWrapper(byte[] payload) {
        super(payload);
        init();
    }

    private void init() {

        byte[] bytes = super.getData();

        // 依次读取每个包，每个包的开头是 0x00 0x00 0x01 PsPlayload
        // 每个ps包的类型不同，长度也不同，所以需要根据类型来判断长度
        // 读取到下一个包的开头时，就可以确定当前包的长度

        for (int i = 0; i < bytes.length - 4; i++) {
            if (bytes[i] != 0x00 || bytes[i + 1] != 0x00 || bytes[i + 2] != 0x01) {
                continue;
            }
            switch (bytes[i + 3]) {

                // PS 包头
                case (byte) 0xBA: {
                    byte[] psHeader = getPsHeader(bytes, i);
                    if (psHeader.length == 0) {
                        continue;
                    }
                    PS_MAP.add(RtpPayload.PsHeaderPack.HEADER, psHeader);
                    i = i + psHeader.length - 1;
                    break;
                }

                // SystemHeader
                case (byte) 0xBB: {
                    byte[] psSysHeader = getSystemHeader(bytes, i);
                    if (psSysHeader.length == 0) {
                        continue;
                    }
                    PS_MAP.add(RtpPayload.PsHeaderPack.SYSTEM_HEADER, psSysHeader);
                    i = i + psSysHeader.length - 1;
                    break;
                }

                // ProgramStreamMap
                case (byte) 0xBC: {
                    byte[] psmHeader = getPsmHeader(bytes, i);
                    if (psmHeader.length == 0) {
                        continue;
                    }
                    PS_MAP.add(RtpPayload.PsHeaderPack.PROGRAM_STREAM_MAP, psmHeader);
                    i = i + psmHeader.length - 1;
                    break;
                }

                // 视频 PES Header
                case (byte) 0xE0: {
                    byte[] pesHeader = getPesHeader(bytes, i);
                    if (pesHeader.length == 0) {
                        continue;
                    }
                    PS_MAP.add(RtpPayload.PsHeaderPack.VIDEO_PES_HEADER, pesHeader);
                    i = i + pesHeader.length - 1;
                    break;
                }

                // 音频 PES Header
                case (byte) 0xC0: {
                    byte[] pesHeader = getPesHeader(bytes, i);
                    if (pesHeader.length == 0) {
                        continue;
                    }
                    PS_MAP.add(RtpPayload.PsHeaderPack.AUDIO_PES_HEADER, pesHeader);
                    i = i + pesHeader.length - 1;
                    break;
                }

                default: {
                    continue;
                }
            }
        }
    }

    public Optional<byte[]> getHeader() {
        return Optional.ofNullable(PS_MAP.getFirst(RtpPayload.PsHeaderPack.HEADER));
    }

    public Optional<byte[]> getSystemHeader() {
        return Optional.ofNullable(PS_MAP.getFirst(RtpPayload.PsHeaderPack.SYSTEM_HEADER));
    }

    public Optional<byte[]> getPsmHeader() {
        return Optional.ofNullable(PS_MAP.getFirst(RtpPayload.PsHeaderPack.PROGRAM_STREAM_MAP));
    }

    public Optional<PsmWrapper> getPsmWrapper() {
        return getPsmHeader().map(PsmWrapper::new);
    }

    public Optional<List<byte[]>> getVideoPesHeader() {
        return Optional.ofNullable(PS_MAP.get(RtpPayload.PsHeaderPack.VIDEO_PES_HEADER));
    }

    public Optional<List<byte[]>> getAudioPesHeader() {
        return Optional.ofNullable(PS_MAP.get(RtpPayload.PsHeaderPack.AUDIO_PES_HEADER));
    }

    private static byte[] getPsHeader(byte[] bytes, int startIndex) {
        // 获取psHeader长度
        byte lenFlag = bytes[startIndex + 13];
        // psHeader扩展字段长度
        int len = lenFlag & 0x07;

        byte[] psHeaders = new byte[13 + len + 1];
        try {
            System.arraycopy(bytes, startIndex, psHeaders, 0, psHeaders.length);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        return psHeaders;
    }

    private static byte[] getSystemHeader(byte[] bytes, int startIndex) {

        // SysHeader字段长度
        int len = ByteUtil.bytesToShort(bytes, startIndex + 4, ByteOrder.BIG_ENDIAN);

        byte[] psSysHeaders = new byte[6 + len];
        System.arraycopy(bytes, startIndex, psSysHeaders, 0, psSysHeaders.length);
        return psSysHeaders;
    }

    private static byte[] getPsmHeader(byte[] bytes, int startIndex) {
        // PSM字段长度
        int len = ByteUtil.bytesToShort(bytes, startIndex + 4, ByteOrder.BIG_ENDIAN);

        byte[] psmHeaders = new byte[6 + len];
        System.arraycopy(bytes, startIndex, psmHeaders, 0, psmHeaders.length);
        return psmHeaders;
    }

    private static byte[] getPesHeader(byte[] bytes, int startIndex) {
        // PES字段长度
        int len = Short.toUnsignedInt(ByteUtil.bytesToShort(bytes, startIndex + 4, ByteOrder.BIG_ENDIAN));

        try {
            byte[] pesHeaders = new byte[6 + len];
            System.arraycopy(bytes, startIndex, pesHeaders, 0, pesHeaders.length);
            return pesHeaders;
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }

}
